#include "trap.h"
#include "../../include/riscv.h"
#include "../boot/printf.h"
#include "../proc/proc.h"
#include "syscall.h"
// 声明汇编函数
extern void kernelvec(void);
extern char trampoline[];          // 汇编全局变量
extern void userret(uint64 satp);  // 汇编全局函数
extern char uservec[];
// 全局变量
volatile uint64 ticks = 0;
interrupt_handler_t irq_table[MAX_IRQ] = {0};

// 测试相关变量
volatile int test_interrupt_count = 0;
volatile int test_mode = 0;  // 0: normal, 1: testing

// 初始化中断系统
void trap_init(void) {
    // 清零中断向量表
    for (int i = 0; i < MAX_IRQ; i++) {
        irq_table[i] = 0;
    }
    
    // 注册默认中断处理函数
    register_interrupt(IRQ_S_TIMER, handle_timer);
    register_interrupt(IRQ_S_EXTERNAL, handle_external);
    register_interrupt(IRQ_S_SOFTWARE, handle_software);
    
    w_stvec((uint64)kernelvec);
    
    // 启用Supervisor模式中断
    intr_on();
    
    printf("Trap system initialized\n");
}

// 注册中断处理函数
void register_interrupt(int irq, interrupt_handler_t handler) {
    if (irq >= 0 && irq < MAX_IRQ) {
        irq_table[irq] = handler;
    }
}

// 注销中断处理函数
void unregister_interrupt(int irq) {
    if (irq >= 0 && irq < MAX_IRQ) {
        irq_table[irq] = 0;
        printf("Unregistered interrupt handler for IRQ %d\n", irq);
    }
}

// 启用特定中断
void enable_interrupt(int irq) {
    uint64 sie = r_sie();
    switch (irq) {
        case IRQ_S_TIMER:
            w_sie(sie | SIE_STIE);
            break;
        case IRQ_S_EXTERNAL:
            w_sie(sie | SIE_SEIE);
            break;
        default:
            printf("Cannot enable IRQ %d - not supported\n", irq);
            break;
    }
}

// 禁用特定中断
void disable_interrupt(int irq) {
    uint64 sie = r_sie();
    switch (irq) {
        case IRQ_S_TIMER:
            w_sie(sie & ~SIE_STIE);
            break;
        case IRQ_S_EXTERNAL:
            w_sie(sie & ~SIE_SEIE);
            break;
        default:
            printf("Cannot disable IRQ %d - not supported\n", irq);
            break;
    }
}

// 内核态陷阱处理函数
void kerneltrap(void) {
    uint64 sepc = r_sepc();
    uint64 scause = r_scause();
    uint64 stval = r_stval();
    printf("kerneltrap\n");
    // 注意：这里没有实现中断嵌套
    // 当前实现是在中断处理期间关闭中断
    // 如果需要支持中断嵌套，需要添加嵌套计数器
    
    // 检查是中断还是异常
    if (scause & (1UL << 63)) {
        // 中断处理
        uint64 irq = scause & 0xFF;
        if (irq < MAX_IRQ && irq_table[irq]) {
            irq_table[irq]();
        } else {
            printf("Unhandled interrupt: %lu\n", irq);
        }
    } else {
        // 异常处理
        handle_exception(scause, sepc, stval);
    }
}

// 时钟中断处理函数
void handle_timer(void) {
    ticks++;
    
    if (test_mode) {
        // 测试模式下增加测试计数
        test_interrupt_count++;
        printf("Test timer interrupt %d occurred\n", test_interrupt_count);
    } 
    // 设置下次时钟中断时间
    sbi_set_timer(get_time() + 1000000);
    
    // 时钟中断可能触发调度
    // 注：实际的调度逻辑应该在进程管理模块中实现
    // 这里只是标记时钟中断已经发生
}

// 外部中断处理函数
void handle_external(void) {
    printf("External interrupt occurred\n");
}

// 软件中断处理函数
void handle_software(void) {
    printf("Software interrupt occurred\n");
}

// 异常处理函数
void handle_exception(uint64 scause, uint64 sepc, uint64 stval) {
    switch (scause) {
        case EXC_ECALL_U_MODE:
            printf("User mode environment call at PC=0x%lx\n", sepc);
            handle_ecall(sepc, stval);
            break;
        case EXC_ECALL_S_MODE:
            printf("Supervisor mode environment call at PC=0x%lx\n", sepc);
            handle_ecall(sepc, stval);
            break;
        case EXC_INSTRUCTION_PAGE_FAULT:
            handle_instruction_page_fault(sepc, stval);
            break;
        case EXC_LOAD_PAGE_FAULT:
            handle_load_page_fault(sepc, stval);
            break;
        case EXC_STORE_PAGE_FAULT:
            handle_store_page_fault(sepc, stval);
            break;
        case EXC_ILLEGAL_INSTRUCTION:
            handle_illegal_instruction(sepc, stval);
            break;
        case EXC_INSTRUCTION_MISALIGNED:
            printf("Instruction address misaligned at PC=0x%lx, address=0x%lx\n", sepc, stval);
            break;
        case EXC_INSTRUCTION_ACCESS_FAULT:
            printf("Instruction access fault at PC=0x%lx, address=0x%lx\n", sepc, stval);
            break;
        case EXC_BREAKPOINT:
            printf("Breakpoint at PC=0x%lx\n", sepc);
            break;
        case EXC_LOAD_MISALIGNED:
            printf("Load address misaligned at PC=0x%lx, address=0x%lx\n", sepc, stval);
            break;
        case EXC_LOAD_ACCESS_FAULT:
            printf("Load access fault at PC=0x%lx, address=0x%lx\n", sepc, stval);
            break;
        case EXC_STORE_MISALIGNED:
            printf("Store address misaligned at PC=0x%lx, address=0x%lx\n", sepc, stval);
            break;
        case EXC_STORE_ACCESS_FAULT:
            printf("Store access fault at PC=0x%lx, address=0x%lx\n", sepc, stval);
            break;
        default:
            printf("Unhandled exception: cause=%d, sepc=0x%lx, stval=0x%lx\n", 
                   scause, sepc, stval);
            break;
    }
}

// 指令页故障处理
void handle_instruction_page_fault(uint64 sepc, uint64 stval) {
    printf("Instruction page fault at PC=0x%lx, address=0x%lx\n", sepc, stval);
    
    // 检查地址是否在有效范围内
    if (stval < 0x80000000 || stval > 0x80000000 + 0x200000) {
        printf("FATAL: Instruction page fault at invalid address\n");
        // 对于测试，我们跳过导致故障的指令
        w_sepc(sepc + 4);
        return;
    }
    
    // 这里应该实现页故障处理逻辑
    // 1. 检查页表项
    // 2. 分配新页面
    // 3. 更新页表
    // 4. 刷新TLB
    printf("Instruction page fault handled - would allocate page\n");
    // 跳过导致故障的指令
    w_sepc(sepc + 4);
}

// 加载页故障处理
void handle_load_page_fault(uint64 sepc, uint64 stval) {
    printf("Load page fault at PC=0x%lx, address=0x%lx\n", sepc, stval);
    
    // 检查地址是否在有效范围内
    if (stval < 0x80000000 || stval > 0x80000000 + 0x8000000) {
        printf("FATAL: Load page fault at invalid address\n");
        panic("Load page fault at invalid address,exiting...");
        return;
    }
    
    // 页故障处理逻辑
   panic("Load page fault \n");
}

// 存储页故障处理
void handle_store_page_fault(uint64 sepc, uint64 stval) {
    printf("Store page fault at PC=0x%lx, address=0x%lx\n", sepc, stval);
    
    // 检查地址是否在有效范围内
    if (stval < 0x80000000 || stval > 0x80000000 + 0x8000000) {
        printf("FATAL: Store page fault at invalid address\n");
        // 对于测试，我们跳过导致故障的指令
        w_sepc(sepc + 4);
        return;
    }
    
    // 页故障处理逻辑
    printf("Store page fault handled - would allocate page\n");
    // 跳过导致故障的指令
    w_sepc(sepc + 4);
}

// 非法指令处理
void handle_illegal_instruction(uint64 sepc, uint64 stval) {
    printf("Illegal instruction at PC=0x%lx, instruction=0x%lx\n", sepc, stval);
    
    // 分析非法指令
    uint32 instruction = (uint32)stval;
    printf("Illegal instruction details:\n");
    printf("  Opcode: 0x%02x\n", instruction & 0x7F);
    printf("  Function3: 0x%x\n", (instruction >> 12) & 0x7);
    printf("  Function7: 0x%02x\n", (instruction >> 25) & 0x7F);
    
    // 跳过导致异常的指令，避免无限循环
    w_sepc(sepc + 4);  // RISC-V指令长度为4字节
    
    printf("Illegal instruction handled, skipping to next instruction\n");
}

// 环境调用处理（系统调用）
void handle_ecall(uint64 sepc, uint64 stval) {
    printf("Environment call at PC=0x%lx\n", sepc);
    
    // 获取系统调用号（通常在a7寄存器中）
    // 这里我们模拟系统调用处理
    printf("System call handling:\n");
    printf("  Call number: (would read from a7 register)\n");
    printf("  Arguments: (would read from a0-a6 registers)\n");
    
    // 系统调用分发逻辑
    // switch (syscall_number) {
    //     case SYS_EXIT: handle_sys_exit(); break;
    //     case SYS_WRITE: handle_sys_write(); break;
    //     default: printf("Unknown system call\n"); break;
    // }
    
    printf("System call handled\n");
}



// ========== 用户态Trap入口 ===========
uint64 usertrap(void) {
    uint64 scause = r_scause();
    uint64 sepc = r_sepc();
    uint64 stval = r_stval();
    //printf("usertrap\n");
    if((r_sstatus() & SSTATUS_SPP) != 0)
    panic("usertrap: not from user mode");
  
    // send interrupts and exceptions to kerneltrap(),
    // since we're now in the kernel.
    w_stvec((uint64)kernelvec);
  
    struct proc *p = myproc();
    
    // save user program counter.
    p->trapframe->epc = sepc;
    
    if(scause == EXC_ECALL_U_MODE){
      // system call
      if(killed(p))
        exit(-1);
  
      // sepc points to the ecall instruction,
      // but we want to return to the next instruction.
      p->trapframe->epc += 4;
  
      // an interrupt will change sepc, scause, and sstatus,
      // so enable only now that we're done with those registers.
      intr_on();
  
      syscall();
    } else if(scause & (1UL << 63)) {
      // 中断处理
      uint64 irq = scause & 0xFF;
      if (irq == IRQ_S_TIMER) {
        // 时钟中断
        handle_timer();
        if(killed(p))
          exit(-1);
        yield();
      } else if (irq < MAX_IRQ && irq_table[irq]) {
        // 其他已注册的中断
        irq_table[irq]();
      } else {
        printf("usertrap(): unexpected interrupt %lu\n", irq);
        setkilled(p);
      }
    } else if((scause == EXC_LOAD_PAGE_FAULT || scause == EXC_STORE_PAGE_FAULT) &&
              vmfault(p->pagetable, stval, (scause == EXC_LOAD_PAGE_FAULT)? 1 : 0) != 0) {
      // page fault on lazily-allocated page
      // vmfault成功处理了页错误
    } else {
      // 其他异常
      printf("usertrap(): unexpected scause 0x%lx pid=%d\n", scause, p->pid);
      printf("            sepc=0x%lx stval=0x%lx\n", sepc, stval);
      setkilled(p);
    }
  
    if(killed(p))
      exit(-1);
  
    prepare_return();

    // the user page table to switch to, for trampoline.S
    uint64 satp = MAKE_SATP(p->pagetable);
  
    // return to trampoline.S; satp value in a0.
    return satp;
}



//
// set up trapframe and control registers for a return to user space
//
void
prepare_return(void)
{
  struct proc *p = myproc();

  // we're about to switch the destination of traps from
  // kerneltrap() to usertrap(). because a trap from kernel
  // code to usertrap would be a disaster, turn off interrupts.
  intr_off();

  // send syscalls, interrupts, and exceptions to uservec in trampoline.S
  uint64 trampoline_uservec = TRAMPOLINE + (uservec - trampoline);
  w_stvec(trampoline_uservec);

  // set up trapframe values that uservec will need when
  // the process next traps into the kernel.
  p->trapframe->kernel_satp = r_satp();         // kernel page table
  p->trapframe->kernel_sp = p->kstack + PGSIZE; // process's kernel stack
  p->trapframe->kernel_trap = (uint64)usertrap;
  p->trapframe->kernel_hartid = r_tp();         // hartid for cpuid()

  // set up the registers that trampoline.S's sret will use
  // to get to user space.
  
  // set S Previous Privilege mode to User.
  unsigned long x = r_sstatus();
  x &= ~SSTATUS_SPP; // clear SPP to 0 for user mode
  x |= SSTATUS_SPIE; // enable interrupts in user mode
  w_sstatus(x);
  
  // set S Exception Program Counter to the saved user pc.
  w_sepc(p->trapframe->epc);
}


// SBI时钟接口实现

// 获取当前时间
uint64 get_time(void) {
    return r_time();
}

// 设置下次时钟中断时间
void sbi_set_timer(uint64 time) {
    w_stimecmp(time);
}

